"""
循环工作流代理示例 - 代码质量优化系统
演示如何使用LoopAgent进行迭代的代码质量改进
"""

from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.tools.function_tool import FunctionTool
from google.adk.tools.tool_context import ToolContext

import os
from google.adk.models.lite_llm import LiteLlm

# API密钥配置
DEFAULT_DASHSCOPE_API_KEY = "sk-f227634bb56*************************232"  # 请替换为您的实际密钥
DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY", DEFAULT_DASHSCOPE_API_KEY)

# 常量定义
APP_NAME = "loop_code_optimizer"
USER_ID = "dev_user_01"
GEMINI_MODEL = LiteLlm(
                model="openai/qwen-turbo",  # 使用通义千问Turbo模型
                api_key=DASHSCOPE_API_KEY,
                api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

# 状态键定义
STATE_ORIGINAL_CODE = "original_code"
STATE_CURRENT_CODE = "current_code"
STATE_ANALYSIS_REPORT = "analysis_report"
STATE_OPTIMIZATION_SUGGESTIONS = "optimization_suggestions"
STATE_QUALITY_SCORE = "quality_score"
STATE_ITERATION_COUNT = "iteration_count"
STATE_OPTIMIZATION_HISTORY = "optimization_history"

# 完成信号
COMPLETION_SIGNAL = "OPTIMIZATION_COMPLETE"

# 工具函数定义
def initialize_code_optimization(tool_context: ToolContext):
    """初始化代码优化的工具"""
    print(f"[Tool] initialize_code_optimization called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 强制防止重复初始化 - 如果已经调用过，直接返回
    call_count = session_state.get("init_call_count", 0)
    session_state["init_call_count"] = call_count + 1
    
    if call_count > 0:
        print(f"[Debug] 防止重复调用 - 第{call_count + 1}次调用，直接返回")
        return {
            "status": "already_initialized", 
            "message": f"已初始化（第{call_count + 1}次调用）",
            "code_length": len(session_state.get(STATE_CURRENT_CODE, ""))
        }
    
    # 尝试从多个来源获取用户输入
    user_input = ""
    
    # 方法1: 从invocation context的user_content获取
    if hasattr(tool_context._invocation_context, 'user_content') and tool_context._invocation_context.user_content:
        user_content = tool_context._invocation_context.user_content
        if hasattr(user_content, 'parts') and user_content.parts:
            # 从Content对象的parts中提取文本
            for part in user_content.parts:
                if hasattr(part, 'text') and part.text:
                    user_input = part.text.strip()
                    break
        elif user_content:
            user_input = str(user_content).strip()
    
    # 方法2: 从session state获取
    if not user_input:
        user_input = session_state.get("user_input", "")
    
    # 方法3: 从session的消息历史获取最新的用户消息
    if not user_input and hasattr(session, 'messages') and session.messages:
        for message in reversed(session.messages):
            if hasattr(message, 'author') and message.author == 'user':
                if hasattr(message, 'content') and message.content:
                    user_input = str(message.content).strip()
                    break
    
    print(f"[Debug] 获取到的用户输入: '{user_input}'")
    
    # 保存用户输入到session state以供后续使用
    if user_input:
        session_state["user_input"] = user_input
    
    current_code = session_state.get(STATE_CURRENT_CODE, "")
    original_code = session_state.get(STATE_ORIGINAL_CODE, "")
    
    # 检查用户输入是否是简单问候或非代码内容
    simple_greetings = ["你好", "hello", "hi", "嗨", "您好", "测试"]
    if user_input and user_input.lower().strip() in [g.lower() for g in simple_greetings]:
        # 对于简单问候，直接返回友好回应并标记为完成
        session_state["simple_greeting"] = True
        session_state[STATE_QUALITY_SCORE] = 10.0  # 设置高分以触发退出
        session_state["initialization_completed"] = True  # 标记初始化完成
        # 设置escalate来跳过后续的优化循环
        tool_context.actions.escalate = True
        print(f"[Debug] 检测到简单问候: '{user_input}', 设置escalate=True")
        return {
            "status": "simple_greeting_detected", 
            "message": "检测到简单问候，无需代码优化",
            "user_input": user_input
        }
    
    # 检查用户输入是否包含代码
    if user_input and (
        "def " in user_input or 
        "class " in user_input or 
        "import " in user_input or
        "from " in user_input or
        "print(" in user_input or
        "return " in user_input or
        len(user_input.split('\n')) > 3  # 多行输入可能是代码
    ):
        # 用户提供了代码，使用用户的代码
        session_state[STATE_ORIGINAL_CODE] = user_input
        session_state[STATE_CURRENT_CODE] = user_input
        session_state["initialization_completed"] = True  # 标记初始化完成
        # 不设置escalate，让流程继续到优化循环
        print(f"[Debug] 检测到用户代码，准备进入优化循环")
        return {"status": "initialized_with_user_code", "code_length": len(user_input)}
    
    # 如果没有代码，使用示例代码
    if not current_code and not original_code:
        sample_code = '''
def calculate_total(items):
    total = 0
    for i in range(len(items)):
        if items[i] > 0:
            total = total + items[i]
        else:
            print("Negative value found")
    return total

def process_data(data):
    result = []
    for item in data:
        if item != None:
            if len(item) > 0:
                result.append(item.upper())
    return result

class DataProcessor:
    def __init__(self):
        self.data = []
    
    def add_item(self, item):
        self.data.append(item)
    
    def get_total(self):
        return calculate_total(self.data)
'''
        session_state[STATE_ORIGINAL_CODE] = sample_code
        session_state[STATE_CURRENT_CODE] = sample_code
        session_state["initialization_completed"] = True  # 标记初始化完成
        # 不设置escalate，让流程继续到优化循环
        print(f"[Debug] 使用示例代码，准备进入优化循环")
        return {"status": "initialized_with_sample", "code_length": len(sample_code)}
    
    # 如果有原始代码但没有当前代码，复制原始代码
    elif original_code and not current_code:
        session_state[STATE_CURRENT_CODE] = original_code
        session_state["initialization_completed"] = True  # 标记初始化完成
        return {"status": "initialized_from_original", "code_length": len(original_code)}
    
    # 如果都有，保持现状
    else:
        session_state["initialization_completed"] = True  # 标记初始化完成
        return {"status": "already_initialized", "code_length": len(current_code)}

def analyze_code_quality(tool_context: ToolContext):
    """分析代码质量的工具"""
    print(f"[Tool] analyze_code_quality called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 检查是否是简单问候
    if session_state.get("simple_greeting", False):
        return {
            "overall_score": 10.0,
            "metrics": {"greeting_response": 10.0},
            "iteration": 1,
            "issues_found": 0,
            "improvements_needed": False,
            "message": "简单问候，无需代码分析"
        }
    
    current_code = session_state.get(STATE_CURRENT_CODE, "")
    iteration = session_state.get(STATE_ITERATION_COUNT, 0)
    
    # 改进的质量评分逻辑 - 更快收敛
    # 模拟代码质量分析
    # 在实际应用中，这里可以集成pylint, flake8, sonarqube等工具
    base_scores = {
        "complexity_score": 6.0,
        "readability_score": 5.5,
        "maintainability_score": 6.2,
        "performance_score": 7.0,
        "security_score": 8.0
    }
    
    # 每次迭代显著提升分数，确保在2-3轮内达到8.0
    improvement_factor = 1.0 + (iteration * 0.4)  # 每轮提升40%
    quality_metrics = {
        key: min(10.0, score * improvement_factor) 
        for key, score in base_scores.items()
    }
    
    overall_score = sum(quality_metrics.values()) / len(quality_metrics)
    
    analysis_report = {
        "overall_score": round(overall_score, 2),
        "metrics": quality_metrics,
        "iteration": iteration + 1,
        "issues_found": max(0, 3 - iteration),  # 减少发现的问题数
        "improvements_needed": overall_score < 8.0
    }
    
    session_state[STATE_ANALYSIS_REPORT] = analysis_report
    session_state[STATE_QUALITY_SCORE] = overall_score
    session_state[STATE_ITERATION_COUNT] = iteration + 1
    
    print(f"[Quality Analysis] Iteration {iteration + 1}, Score: {overall_score:.2f}")
    
    return analysis_report

def exit_optimization_loop(tool_context: ToolContext):
    """退出优化循环的工具"""
    print(f"[Tool] exit_optimization_loop called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 检查是否是简单问候
    if session_state.get("simple_greeting", False):
        user_input = session_state.get("user_input", "")
        tool_context.actions.escalate = True
        return {
            "status": "greeting_completed", 
            "message": f"已回应用户问候：{user_input}",
            "final_score": 10.0
        }
    
    # 记录优化历史
    history = session_state.get(STATE_OPTIMIZATION_HISTORY, [])
    final_score = session_state.get(STATE_QUALITY_SCORE, 0)
    iteration = session_state.get(STATE_ITERATION_COUNT, 0)
    
    completion_reason = "Quality threshold reached"
    if iteration >= 3:
        completion_reason = "Maximum iterations reached"
    elif final_score >= 8.0:
        completion_reason = "Quality threshold reached"
    
    history.append({
        "final_iteration": iteration,
        "final_score": final_score,
        "completion_reason": completion_reason
    })
    
    session_state[STATE_OPTIMIZATION_HISTORY] = history
    tool_context.actions.escalate = True
    
    print(f"[Exit Loop] Reason: {completion_reason}, Final Score: {final_score:.2f}")
    
    return {"status": "optimization_completed", "final_score": final_score, "reason": completion_reason}

def save_optimization_step(tool_context: ToolContext):
    """保存优化步骤的工具"""
    print(f"[Tool] save_optimization_step called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    history = session_state.get(STATE_OPTIMIZATION_HISTORY, [])
    current_code = session_state.get(STATE_CURRENT_CODE, "")
    suggestions = session_state.get(STATE_OPTIMIZATION_SUGGESTIONS, "")
    score = session_state.get(STATE_QUALITY_SCORE, 0)
    iteration = session_state.get(STATE_ITERATION_COUNT, 0)
    
    step_record = {
        "iteration": iteration,
        "quality_score": score,
        "suggestions_applied": suggestions[:100] + "..." if len(suggestions) > 100 else suggestions,
        "code_length": len(current_code),
        "timestamp": f"Step {iteration}"
    }
    
    history.append(step_record)
    session_state[STATE_OPTIMIZATION_HISTORY] = history
    
    return step_record


analyze_code_quality_tool = FunctionTool(analyze_code_quality)
initialize_code_optimization_tool = FunctionTool(initialize_code_optimization)

# 代理定义

# 步骤1: 代码分析代理
code_analyzer_agent = LlmAgent(
    name="CodeAnalyzerAgent",
    model=GEMINI_MODEL,
    instruction=f"""你是一个专业的代码质量分析专家。

    请分析当前会话状态中的代码质量。
    
    请调用analyze_code_quality工具来分析代码质量，包括：
    - 代码复杂度
    - 可读性
    - 可维护性
    - 性能
    - 安全性
    
    分析完成后，请简要说明发现的主要问题和当前的质量评分。
    """,
    description="分析代码质量，识别需要改进的方面",
    tools=[analyze_code_quality_tool],
    include_contents='none'
)

# 步骤2: 优化建议代理
optimization_advisor_agent = LlmAgent(
    name="OptimizationAdvisorAgent",
    model=GEMINI_MODEL,
    instruction=f"""你是一个代码优化顾问。

    基于代码质量分析报告，提供具体的优化建议。
    请从会话状态中获取分析报告和当前代码信息。
    
    请提供以下方面的优化建议：
    1. 代码结构改进
    2. 性能优化建议
    3. 可读性提升
    4. 最佳实践应用
    5. 潜在bug修复
    
    建议应该：
    - 具体可操作
    - 优先级明确
    - 包含代码示例
    - 解释改进原因
    
    如果质量评分已经达到8.0以上，请在建议末尾添加"质量已达标，建议完成优化。"
    """,
    description="基于分析结果提供具体的代码优化建议",
    output_key=STATE_OPTIMIZATION_SUGGESTIONS,
    include_contents='none'
)

save_optimization_step_tool = FunctionTool(save_optimization_step)
# 步骤3: 代码重构代理
code_refactor_agent = LlmAgent(
    name="CodeRefactorAgent",
    model=GEMINI_MODEL,
    instruction=f"""你是一个代码重构专家。

    根据优化建议重构代码。
    请从会话状态中获取当前代码和优化建议。
    
    请根据建议重构代码，重点关注：
    1. 应用建议的改进措施
    2. 保持代码功能不变
    3. 提高代码质量
    4. 遵循编程最佳实践
    
    请先调用save_optimization_step工具保存当前优化步骤，然后输出重构后的完整代码。
    
    输出格式：
    ```python
    # 重构后的代码
    [重构后的完整代码]
    ```
    """,
    description="根据优化建议重构代码",
    output_key=STATE_CURRENT_CODE,
    tools=[save_optimization_step_tool],
    include_contents='none'
)

exit_optimization_loop_tool = FunctionTool(exit_optimization_loop)
# 步骤4: 质量检查代理
quality_checker_agent = LlmAgent(
    name="QualityCheckerAgent",
    model=GEMINI_MODEL,
    instruction=f"""你是一个代码质量检查员。

    检查优化结果并决定是否继续迭代。
    请从会话状态中获取分析报告、优化建议和当前质量评分。
    
    评估标准：
    1. 质量评分是否达到8.0以上
    2. 是否还有明显的改进空间
    3. 优化建议是否提到"质量已达标"
    4. 是否是简单问候（无需优化）
    
    决策规则：
    - 如果是简单问候（session中有simple_greeting=True），立即调用exit_optimization_loop工具
    - 如果质量评分 >= 8.0，调用exit_optimization_loop工具结束优化
    - 如果建议中包含"质量已达标"，调用exit_optimization_loop工具结束优化
    - 如果迭代次数 >= 3，调用exit_optimization_loop工具结束优化（防止无限循环）
    - 否则，说明需要继续下一轮优化
    
    重要：必须明确调用exit_optimization_loop工具来结束循环，不能只是说明需要结束。
    
    请根据当前情况做出决策。
    """,
    description="检查优化质量，决定是否继续迭代",
    tools=[exit_optimization_loop_tool],
    include_contents='none'
)

# 创建循环工作流代理
optimization_loop = LoopAgent(
    name="OptimizationLoop",
    description="代码质量优化循环 - 重复执行分析、建议、重构、检查直到质量达标",
    sub_agents=[
        code_analyzer_agent,
        optimization_advisor_agent,
        code_refactor_agent,
        quality_checker_agent
    ],
    max_iterations=5  # 最多5轮优化
)

# 初始化代理
initialization_agent = LlmAgent(
    name="InitializationAgent",
    model=GEMINI_MODEL,
    instruction="""你是初始化专家。

执行步骤：
1. 调用 initialize_code_optimization 工具
2. 根据工具返回的 status 字段回应：
   - "simple_greeting_detected": 回复"您好！"
   - "initialized_with_user_code": 回复"代码已接收，开始分析。"
   - "initialized_with_sample": 回复"使用示例代码演示。"
   - "already_initialized": 回复"系统已就绪。"
3. 最后说"初始化完成"

只调用工具一次，然后结束。""",
    description="初始化代码优化流程",
    tools=[initialize_code_optimization_tool]
)

# 创建完整的工作流
root_agent = SequentialAgent(
    name=APP_NAME,
    description="代码质量优化系统 - 使用循环工作流进行迭代优化",
    sub_agents=[
        initialization_agent,
        optimization_loop
    ]
)

# 用于测试的辅助函数
def create_test_session_with_code(session_service, code_content: str):
    """创建包含测试代码的会话"""
    session = session_service.create_session(APP_NAME, USER_ID)
    session.state[STATE_ORIGINAL_CODE] = code_content
    session.state[STATE_CURRENT_CODE] = code_content
    session.state[STATE_ITERATION_COUNT] = 0
    session.state[STATE_OPTIMIZATION_HISTORY] = []
    return session

# 示例代码内容
SAMPLE_CODE = '''
def calculate_total(items):
    total = 0
    for i in range(len(items)):
        if items[i] > 0:
            total = total + items[i]
        else:
            print("Negative value found")
    return total

def process_data(data):
    result = []
    for item in data:
        if item != None:
            if len(item) > 0:
                result.append(item.upper())
    return result

class DataProcessor:
    def __init__(self):
        self.data = []
    
    def add_item(self, item):
        self.data.append(item)
    
    def get_total(self):
        return calculate_total(self.data)
''' 